package com.afdor.rws.dlock.jedis.sharded;

import com.afdor.rws.annotation.extension.Extension;
import com.afdor.rws.dlock.lock.ReentrantLock;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.ShardedJedis;
import redis.clients.jedis.ShardedJedisPool;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.LockSupport;

/**
 * 采用jedis客户端实现的基于redis的分布式锁，该锁效率高，但存在可用性问题，不能用于要求高稳定性环境或其他敏感业务场景，请谨慎使用
 *
 * @author 悭梵
 * @date Created in 2018-09-04 14:59
 */
@Slf4j
@ToString
@EqualsAndHashCode
@Extension(value = "shardedJedis", scope = Extension.Scope.PROTOTYPE, interfaceClass = ReentrantLock.class)
public class DistributedReentrantLock implements ReentrantLock {

    private static final String ROOT_PATH = "/Distributed/ReentrantLock/";
    private final ConcurrentMap<Thread, LockData> threadData = new ConcurrentHashMap<>();

    /**
     * 重试等待时间
     */
    @Getter
    private int retryAwait = 300;
    @Getter
    private int lockTimeout = 2000;
    @Getter
    protected String lockId;
    @Getter
    protected final ShardedJedisPool shardedJedisPool;

    public DistributedReentrantLock(ShardedJedisPool shardedJedisPool, String lockId) {
        if (StringUtils.isBlank(lockId)) {
            throw new NullPointerException("锁标识为空");
        }
        if (shardedJedisPool == null) {
            throw new NullPointerException("Jedis客户端为空");
        }
        this.lockId = ROOT_PATH + lockId;
        this.shardedJedisPool = shardedJedisPool;
    }

    @Override
    public void lock() {
        tryLock(Integer.MAX_VALUE, TimeUnit.MILLISECONDS);
    }

    @Override
    public boolean tryLock(long timeout, TimeUnit unit) {
        Thread currentThread = Thread.currentThread();
        LockData lockData = threadData.get(currentThread);
        if (lockData != null) {
            lockData.lockCount.incrementAndGet();
            return true;
        }
        String lockVal = tryRedisLock(lockId, timeout, unit);
        if (lockVal != null) {
            threadData.putIfAbsent(currentThread, new LockData(currentThread, lockVal));
            return true;
        }
        return false;
    }

    @Override
    public void unlock() {
        Thread currentThread = Thread.currentThread();
        LockData lockData = threadData.get(currentThread);
        if (lockData == null) {
            throw new IllegalMonitorStateException("You do not own the lock: " + lockId);
        }
        int newLockCount = lockData.lockCount.decrementAndGet();
        if (newLockCount > 0) {
            return;
        }
        if (newLockCount < 0) {
            throw new IllegalMonitorStateException("Lock count has gone negative for lock: " + lockId);
        }
        try {
            unlockRedisLock(lockId, lockData.lockVal);
        } finally {
            threadData.remove(currentThread);
        }
    }

    protected String tryRedisLock(String lockId, long timeout, TimeUnit unit) {
        final long startMillis = System.currentTimeMillis();
        final Long millisToWait = (unit != null) ? unit.toMillis(timeout) : null;
        String lockValue = null;
        while (lockValue == null) {
            lockValue = createRedisKey(lockId);
            if (lockValue != null) {
                break;
            }
            if (System.currentTimeMillis() - startMillis - retryAwait > millisToWait) {
                break;
            }
            LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(retryAwait));
        }
        return lockValue;
    }

    protected String createRedisKey(String lockId) {
        ShardedJedis shardedJedis = null;
        try {
            shardedJedis = shardedJedisPool.getResource();
            Jedis jedis = shardedJedis.getShard(lockId);

            String value = lockId + "_" + ThreadLocalRandom.current().nextDouble(99999, 999999999);

            String luaScript = ""
                    + "\nlocal r = tonumber(redis.call('SETNX', KEYS[1], ARGV[1]));"
                    + "\nredis.call('PEXPIRE', KEYS[1], ARGV[2]);"
                    + "\nreturn r";
            List<String> keys = new ArrayList<String>();
            keys.add(lockId);
            List<String> args = new ArrayList<String>();
            args.add(value);
            args.add(lockTimeout + "");

            Long ret = (Long) jedis.eval(luaScript, keys, args);
            if (ret != null && ret.longValue() == 1L) {
                return value;
            }
        } finally {
            if (shardedJedis != null) {
                shardedJedis.close();
            }
        }
        return null;
    }

    protected void unlockRedisLock(String key, String value) {
        ShardedJedis shardedJedis = null;
        try {
            shardedJedis = shardedJedisPool.getResource();
            Jedis jedis = shardedJedis.getShard(lockId);

            String luaScript = ""
                    + "\nlocal v = redis.call('GET', KEYS[1]);"
                    + "\nlocal r= 0;"
                    + "\nif v == ARGV[1] then"
                    + "\nr =redis.call('DEL', KEYS[1]);"
                    + "\nend"
                    + "\nreturn r";

            List<String> keys = new ArrayList<String>();
            keys.add(key);
            List<String> args = new ArrayList<String>();
            args.add(value);

            jedis.eval(luaScript, keys, args);
        } finally {
            if (shardedJedis != null) {
                shardedJedis.close();
            }
        }
    }


    /**
     * 锁数据
     */
    private static class LockData {
        final String lockVal;
        final Thread owningThread;
        final AtomicInteger lockCount = new AtomicInteger(1);

        private LockData(Thread owningThread, String lockVal) {
            this.owningThread = owningThread;
            this.lockVal = lockVal;
        }
    }
}
